import cv2
import numpy as np
import os
import tqdm


def histogram_matching(source, template):
    """
    Adjust the pixel values of the source image such that its histogram matches that of the template image.
    """
    oldshape = source.shape
    source = source.ravel()
    template = template.ravel()

    # Get the set of unique pixel values and their corresponding indices and counts
    s_values, bin_idx, s_counts = np.unique(
        source, return_inverse=True, return_counts=True
    )
    t_values, t_counts = np.unique(template, return_counts=True)

    # Calculate the empirical cumulative distribution functions (CDF) of the source and template images
    s_quantiles = np.cumsum(s_counts).astype(np.float64) / source.size
    t_quantiles = np.cumsum(t_counts).astype(np.float64) / template.size

    # Interpolate to find the pixel values in the template image that correspond to the quantiles in the source image
    interp_t_values = np.interp(s_quantiles, t_quantiles, t_values)

    return interp_t_values[bin_idx].reshape(oldshape)


# # 读取DEM数据
# dem_image = cv2.imread(file_path, cv2.IMREAD_UNCHANGED)
# dem_image = dem_image[3600:3700, 700:800]
# dem_image = (dem_image - dem_image.min()) / (dem_image.max() - dem_image.min())
# dem_image = dem_image**1 * 255

# dem_image = dem_image.astype("uint8")
# # dem_image = cv2.equalizeHist(dem_image)
# cv2.imwrite("dem.png", dem_image)

# grad_x = cv2.Sobel(dem_image, cv2.CV_64F, 1, 0, ksize=3)
# grad_y = cv2.Sobel(dem_image, cv2.CV_64F, 0, 1, ksize=3)

# grad = cv2.magnitude(grad_x, grad_y)
# orientation = cv2.phase(grad_x, grad_y, angleInDegrees=True)
# grad = (grad - grad.min()) / (grad.max() - grad.min()) * 255
# cv2.imwrite("grad.png", grad)
file_path = "/disk527/DataDisk/a804_cbf/datasets/lunar_crater/lunar_lronac_haworth_sfs-dem_shade_1m_v3.jpg"
file_dir = "/disk527/sdb1/a804_cbf/datasets/haworth/images"
mask_path = "/disk527/sdb1/a804_cbf/datasets/Lunar_Crater_Detection_Data/LRO_DATA/train/M1266003170_mosaic_train_small.png"
resolution = 400
file_dir = os.path.join(file_dir, f"{resolution}")
if not os.path.exists(file_dir):
    os.makedirs(file_dir)
img = cv2.imread(file_path, cv2.IMREAD_UNCHANGED)
for i in range(img.shape[0] // resolution):
    for j in range(img.shape[1] // resolution):
        shade_image = img[
            i * resolution : (i + 1) * resolution,
            j * resolution : (j + 1) * resolution,
        ]
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        # shade_image = (shade_image - shade_image.min()) / (
        #     shade_image.max() - shade_image.min()
        # )
        # shade_image = shade_image * 255
        shade_image = shade_image.astype("uint8")
        # shade_image = cv2.equalizeHist(shade_image)
        shade_image = histogram_matching(shade_image, mask)
        # matched_image = shade_image**2.2 * 255
        shade_image = (shade_image - shade_image.min()) / (
            shade_image.max() - shade_image.min()
        )
        shade_image = shade_image**0.9 * 255
        cv2.imwrite(os.path.join(file_dir, f"{i}_{j}.png"), shade_image)
